from major.models_manager import embedding_model
import numpy as np

model = embedding_model.get_model()


reference = '我爱自然语言学习'
# 候选序列（分词后的列表）
candidates = ['我喜欢自然语言','我不爱自然语言']

reference_embedding = model.embed_query(reference)
candidate_embeddings = model.embed_documents(candidates)

def cos_sim(a, b):
    """计算两个向量的余弦相似度"""
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

for i,candidate in enumerate(candidate_embeddings):
    print(f"真实句子: {reference}")
    print(f"候选序列 {i+1}: {candidates[i]}")
    print(f"相似度: {cos_sim(reference_embedding, candidate)}")
    print("-"*20)

